Skip to content

(2/3) refactor(session): assemble training samples on the session server; records never leave it#1759

Draft
guapisolo wants to merge 1 commit into
mainfrom
refactor/session-samples-op
Draft

(2/3) refactor(session): assemble training samples on the session server; records never leave it#1759
guapisolo wants to merge 1 commit into
mainfrom
refactor/session-samples-op

Conversation

@guapisolo

@guapisolo guapisolo commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Depends on: #1758
Stacked on the parent PR's branch. Review / merge the parent first.

Summary

Assemble training samples on the session server; records never leave the owning instance.

Motivation

Each record's R3 buffer covers the full prefix, so a session's records dump grows quadratically with turns — measured ~3.15 GiB per full GET /sessions/{id} at production shape (an mp-stack figure; the magnitude carries to this stack) — then the rollout driver parsed every byte of it on one interpreter per session. The only consumer of records is Sample assembly, whose inputs all already live on the owning instance. What stays the same: the assembly pipeline (compute, truncate, merge) runs unchanged, with the assembled values locked by golden tests.

Before / After

  • Before / After: the driver fetched the full records dump per session; now POST /sessions/{session_id}/samples runs compute → truncate → merge on the owning instance, synchronously on its event loop (the lock-free get_session invariant).
  • The reply is one safetensors payload whose contract is the single declarative table SAMPLES_VALUE_SPEC (field → codec/dtype): per-sample tensors named {field}.{index} (loss_mask crosses as uint8), json-codec fields riding as the rank-one uint8 tensor _samples_meta.
  • What moved where: the wire codec is born as miles/rollout/session/samples/codec.py, depending only on Sample, NumPy, safetensors; SessionCore.collect_samples, the route, http_utils.post_bytes_no_retry, plus the client cutover land together.
  • agentic_tool_call.generate shrinks to a shell: agent call → collect_samplesempty_reason mapping → today's metadata application order.

Behavior Preservation

  • How we know: golden tests assert the exact assembled Sample field values (per-turn, merged, truncated) derivable from a fixed two-turn records fixture.
  • Only the ten COMPUTED_FIELDS cross the wire on blank templates; the driver overlays them onto deepcopies of its input sample; every other Sample field never crosses — the overlay keeps the driver deepcopy's value verbatim.
  • Three deliberate client deltas, each locked by a test: a non-2xx raises with the server's body text (a 422 carries the assembly assertion verbatim); a collect timeout raises instead of silently ABORTing the sample (a measured data-loss bug in the records path); the session DELETE is attempted on every path.

Verification

  • tests/fast/router/test_session_samples_op.py — golden field values; 422 / 404; empty_reason discriminators; route registered before the catch-all proxy.
  • tests/fast/rollout/session/ — assembly goldens plus codec round-trip / malformed-payload tests; client deltas in test_openai_endpoint_utils.py.
  • 75 passed at this tip across router op, session, client.

Review Focus

  • Scrutinize collect_samples in core.py: assembly must stay synchronous on the loop — offloading to an executor without snapshotting records breaks the lock-free read.
  • Scrutinize the COMPUTED_FIELDS allowlist in samples/codec.py: a computed field missing from it silently keeps the driver's stale value at training time.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Comment thread miles/rollout/session/samples/codec.py Outdated
@@ -0,0 +1,222 @@
"""Black-box sample packing for the `POST /sessions/{id}/samples` reply; only the

@guapisolo guapisolo Jul 22, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a vibed file, but verified by e2e CIs. No tito mismatch rate regression.

@guapisolo
guapisolo force-pushed the refactor/session-samples-move branch from 0d64438 to 8c55e0e Compare July 23, 2026 19:40
@guapisolo
guapisolo force-pushed the refactor/session-samples-op branch from 95a27be to 8fdc5d5 Compare July 23, 2026 19:40
@guapisolo
guapisolo marked this pull request as draft July 24, 2026 20:39
@guapisolo
guapisolo force-pushed the refactor/session-samples-op branch from 8fdc5d5 to a70ff78 Compare July 24, 2026 20:59
@guapisolo
guapisolo force-pushed the refactor/session-samples-move branch from 8c55e0e to b82b089 Compare July 24, 2026 21:13
Base automatically changed from refactor/session-samples-move to main July 24, 2026 21:16
…ds never leave it

POST /sessions/{session_id}/samples (registered before the catch-all proxy, which would otherwise forward it to the inference backend): the owning instance runs compute -> truncate -> merge synchronously on its event loop (the lock-free get_session invariant) and replies with one safetensors payload — per-sample tensors named {field}.{index} for the array fields (loss_mask crosses as uint8, one byte per token), with the json-codec fields riding as the rank-one uint8 tensor _samples_meta (safetensors.numpy.load exposes no header metadata and __metadata__ is reserved by the format), so malformed payloads fail inside safetensors' validated parser instead of hand-rolled framing checks. The wire contract is one declarative table, SAMPLES_VALUE_SPEC: every computed field maps to a ValueSpec naming its codec ("tensor" | "tensor_list" | "json") and pinned wire dtype, and the COMPUTED_FIELDS allowlist derives from the table. Only those fields cross the wire on blank templates; every other Sample field never crosses — the driver overlays the wire fields onto deepcopies of its input sample and keeps its local value verbatim for the rest. Non-strict dtype conversions must be value-exact (a negative loss-mask entry wrapping into uint8 raises instead of shipping corrupt masks). The codec lives in samples/codec.py, depending only on Sample, NumPy, and safetensors (>=0.8.0, now an explicit dependency); deterministic assembly failures return 422 with the assertion text; empty replies carry a no_records/all_truncated discriminator preserving today's ABORTED semantics.

The old records path (GET the full dump, GiB-scale under R3 with per-turn full-prefix arrays, parsed on the driver's single interpreter) is retired from the training path. collect_samples posts once via the new http_utils.post_bytes_no_retry (post() force-decodes json/text and blind-retries; a 5xx here means the owning instance died with the records, a 422 is deterministic), raises on non-2xx with the body text, raises on timeout instead of silently ABORTing the sample (a measured data-loss bug in the records path), and attempts the session DELETE on every path so failed sessions cannot accumulate. agentic_tool_call.generate shrinks to a shell: agent call -> collect_samples -> empty_reason mapping -> metadata application in today's order.

Tests assert golden Sample field values derived from the records fixture (per-turn, merged, truncated), the wire-codec round-trip, malformed-payload contracts, the 422/404/empty_reason/route-order contracts, and the client's behavior deltas.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@guapisolo
guapisolo force-pushed the refactor/session-samples-op branch from a70ff78 to c95ee80 Compare July 24, 2026 21:24
guapisolo added a commit that referenced this pull request Jul 25, 2026
…server v2

The v5 multi-lineage design lands as an OPT-IN second implementation
instead of replacing v1: --use-session-server grows an optional value
(bare flag or "v1" = the untouched linear server, "v2" = tree serving),
and the whole tree stack lives in miles/rollout/session/v2/. The v1
modules stay byte-identical to their pre-branch state.

v2 package (carried from the v5 line, imports adjusted only):
- session_state: always-branch serving over the forest — deepest attach
  point, suffix becomes the branch delta, non-extensions grow siblings or
  new roots; retry semantics moves entirely to the samples-op picker.
  Branch suffixes may carry client assistants (compaction); snapshot
  splice with zero-inheritance-root fallback on canonical-prefix
  divergence. Truncated-path extension is 409 (TruncatedGenerationError,
  defined here — v1 never raises it). --session-strict-append-only is the
  single-chain guard: branch shapes fail loud with attach diagnostics.
- assembly: per-leaf fold (compute -> truncate -> fold), default pick
  (temporal-supersession retry trim; roots never trimmed; non-retry-shaped
  trees 422) and default merge (exactly-once completion masking over the
  surviving set, rewards keyed by response id, fold < agent < server
  metadata layering); both replaceable via --session-sample-picker-path /
  --session-merge-function-path (sync-only, loaded in-process).
- core: v2 twin of SessionCore reusing v1's HTTP plumbing unchanged.

Shared seams, each inert for v1 by construction:
- sessions.py dispatches registry/core on the flag and forwards the
  collect body's "metadata" (agent semantic layer) only under v2.
- codec: encode/decode parameterized by a fields tuple; v1 default keeps
  the wire byte-identical, v2 adds reward + per-sample metadata
  (COMPUTED_FIELDS_V2) with conditional overlay semantics on decode.
- arguments: flag upgrade (nargs="?"), unknown values rejected, and the
  three v2 flags require --use-session-server v2.

Tests: v2 HTTP matrix (tree pins incl. deep/root divergence branching,
strict guard, truncation + compaction), session_state unit matrix, and
the samples-op suite (golden, multi-leaf trim/masking/rewards, hook
lanes) — all against a v2-flagged server; the restored v1 suites and the
byte-exact A-list (git diff against the pre-branch base is empty) pin
that the default path did not move.

Rebased onto the tito/session PR stack (#1628/#1777/#1778 +
#1759/#1760), with the stack's landed shape taken as authoritative:

- supports_midpath_assistant_rerender (True on the base class) moves in
  here with its consumer, the v2 suffix gate. No family sets it False
  anymore: the vendored DeepSeek-V3.2 encoder renders
  position-independently under drop_thinking=False, so the gate text no
  longer names it; a stub family pins the fail-loud branch.
- v2 serving fills SessionRecord.request_timestamp like v1, so folded
  samples keep their lifecycle req_ts across the wire.
- codec: COMPUTED_FIELDS_V2 is re-expressed over the ValueSpec table
  (SAMPLES_VALUE_SPEC_V2 superset), same wire semantics.
- v2 unit tests pin allowed_append_roles=["tool"] explicitly; the ctor
  default is tool+user since the append-roles PR.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant